# !/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : textrank-idf模型.py
# @Author: dongguangwen
# @Date  : 2024-07-14 0:04
import jieba
import networkx as nx
from collections import defaultdict
from math import log

# 示例文本列表
documents = [
    "TextRank是一种基于图的排序模型，用于文本处理。",
    "它基于PageRank算法，并可以提取文本中的关键词和句子。",
    "TextRank算法常用于文本摘要和文档聚类等任务。"
]


# 计算IDF值
def compute_idf(documents):
    idf_dict = {}
    N = len(documents)
    for document in documents:
        words = jieba.lcut(document)
        word_set = set(words)
        for word in word_set:
            if word not in idf_dict:
                idf_dict[word] = 0
            idf_dict[word] += 1
    for word, df in idf_dict.items():
        idf_dict[word] = log(N / (1 + df))  # 使用平滑的IDF计算方式
    return idf_dict


# 构建图并应用TextRank-IDF算法
def textrank_idf(documents, idf_dict, topK=5, max_iter=1000, tol=1.0e-6):
    word_pairs = defaultdict(list)
    for document in documents:
        words = jieba.lcut(document)
        for i, word in enumerate(words):
            for j in range(i + 1, len(words)):
                if words[j] not in idf_dict:
                    continue
                word_pairs[word].append((words[j], idf_dict[words[j]]))
                word_pairs[words[j]].append((word, idf_dict[word]))

    graph = nx.Graph()
    for word, pairs in word_pairs.items():
        for pair, weight in pairs:
            graph.add_edge(word, pair, weight=weight)

    scores = nx.pagerank(graph, weight='weight', max_iter=max_iter, tol=tol)

    # 提取关键词
    keywords = sorted(scores, key=scores.get, reverse=True)[:topK]
    return keywords


# 计算IDF值
idf_dict = compute_idf(documents)

# 应用TextRank-IDF算法提取关键词
keywords = textrank_idf(documents, idf_dict)

print("提取的关键词：", keywords)

"""
提取的关键词： ['是', '一种', '常用', '于', '摘要']
"""